#!/usr/bin/env python

import cPickle as pickle
import multiprocessing
import optparse
import os
import os.path
import sys

import tables

import charmicat

def draw(analysis_name, path):
    analysis = __import__(analysis_name)

    try:
        fp = open(os.path.join(path, 'hists.hpkl'), 'rb')
    except IOError:
        sys.stderr.write("`hists.hpkl' not found in {0}\n".format(path))
        sys.stderr.write("Skipping...\n")
        return 

        
    h = pickle.load(fp)              
    fp.close()

    draw_opts = analysis.cuts.hists[h['total'].name].draw_opts
    draw_opts['hold'] = True
    draw_opts['facecolor'] = '#0066FF'
    h['total'].draw(**draw_opts)
    draw_opts['facecolor'] = '#FF0099'
    draw_opts['hold'] = False
    draw_opts['same'] = True
    h['pass'].draw(**draw_opts)
    

def main():
    op = optparse.OptionParser(usage='%prog [opts] analysis_package [path]')

    op.add_option("-c", "--cutset", dest="cutsets", action='append', default=[],
                  help="Draw histograms for cutset NAME (can be used more than once)")
    op.add_option("-r", "--replace", action="append", dest="replace", default=[], nargs=2,
                  help="Replace in cutset cut named CUT1 with cut named CUT2")
    op.add_option("-s", dest="list", action='append', default=[],
                  help="Draw histograms for the cut NAME (can be used more than once)")
    op.add_option("--hook", dest="hook", default=None, 
                  help="Replace draw with custom function NAME from `analysis_package'.cuts",
                  metavar="NAME")
    op.add_option("--start-on", dest="start_on", default=None,
                  help="Fill hists only for cuts downstream (inclusive) of NAME", metavar='NAME')

    opts, args = op.parse_args()

    if len(args) < 1:
        sys.stderr.write("Please supply the name of an analysis package.\n")
        sys.stderr.write("See `draw_cut_hists.py -h' for more.\n")
        sys.stderr.flush()
        return -1

    analysis_name = args.pop(0)

    try:
        analysis      = __import__(analysis_name)
    except ImportError:
        sys.stderr.write("No module `{0}' found. Please check your $PYTHONPATH.\n".format(analysis_name))
        sys.stderr.write("See `apply_cuts.py -h' for more.\n")
        sys.stderr.flush()
        return -2
            
    if len(args) == 0:
        og_dir = os.getcwd()
    elif len(args) == 1:
        og_dir = os.path.abspath(args.pop(0))
    else:
        sys.stderr.write("Too many command-line arguments--see `draw_cut_hists.py -h' for more.\n")
        sys.stderr.flush()
        return -2
        
    if opts.list or opts.cutsets or opts.replace:
        try:
            cut_list = charmicat.build_cut_list(analysis, opts)
        except optparse.OptionError:
            return -4
        
        cut_dirs = {}
        cut_dir  = og_dir
        start_on_reached = False
        for cut in cut_list:
            cut_dir = os.path.join(cut_dir, cut.name)
            
            if opts.start_on and not start_on_reached:
                if cut.name == opts.start_on:
                    start_on_reached = True
                else:
                    continue
            
            cut_dirs[cut.name] = cut_dir
        
        
    else:
        cut_list = []
        cut_dirs = {}
        for root, dirs, files in os.walk(og_dir):
            for idir in dirs:
                if idir in analysis.cuts.all_dict:
                    cut_list.append(analysis.cuts.all_dict[idir])
                    cut_dirs[idir] = os.path.join(root, idir)

    if opts.hook:
        draw_fcn = getattr(analysis.cuts, opts.hook)                
    else:
        draw_fcn = draw

    for key in cut_dirs:
        if key not in analysis.cuts.hists:
            continue

        try:
            draw_fcn(analysis_name, cut_dirs[key])
        except ValueError, e:
            raise ValueError("Error drawing histogram {0}: {1}".format(key, e))

    return 0

if __name__ == "__main__":
    sys.exit(main())
